"""
Whisper Large V3 Turbo 推理引擎
"""
import torch
import logging
from typing import Dict, List, Optional, Union
import numpy as np
from transformers import pipeline

logger = logging.getLogger(__name__)


class WhisperInference:
    """Whisper推理引擎"""
    
    def __init__(self, model, processor, device: str):
        self.model = model
        self.processor = processor
        self.device = device
        self.pipe = None
        self._init_pipeline()
    
    def _init_pipeline(self):
        """初始化pipeline"""
        try:
            torch_dtype = torch.float16 if self.device != "cpu" else torch.float32
            
            self.pipe = pipeline(
                "automatic-speech-recognition",
                model=self.model,
                tokenizer=self.processor.tokenizer,
                feature_extractor=self.processor.feature_extractor,
                torch_dtype=torch_dtype,
                device=self.device,
            )
            logger.info("Pipeline初始化成功")
        except Exception as e:
            logger.error(f"Pipeline初始化失败: {str(e)}")
            raise
    
    def transcribe(
        self,
        audio: Union[str, np.ndarray, Dict],
        language: Optional[str] = None,
        task: str = "transcribe",
        return_timestamps: bool = False,
        temperature: Union[float, tuple] = (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
        **kwargs
    ) -> Dict:
        """
        转录音频
        
        Args:
            audio: 音频文件路径或音频数组
            language: 语言代码（可选）
            task: 任务类型 ("transcribe" 或 "translate")
            return_timestamps: 是否返回时间戳
            temperature: 采样温度
            **kwargs: 其他生成参数
        
        Returns:
            转录结果字典
        """
        try:
            # max_new_tokens 需要小于 max_target_positions (448)
            # 考虑到 decoder_input_ids 已经包含特殊令牌（通常为 4），设置为 444 以确保不超过限制
            generate_kwargs = {
                "max_new_tokens": 444,
                "num_beams": 1,
                "condition_on_prev_tokens": False,
                "compression_ratio_threshold": 1.35,
                "temperature": temperature,
                "logprob_threshold": -1.0,
                "no_speech_threshold": 0.6,
                "return_timestamps": return_timestamps,
                **kwargs
            }
            
            if language:
                generate_kwargs["language"] = language
            
            if task == "translate":
                generate_kwargs["task"] = "translate"
            
            result = self.pipe(audio, generate_kwargs=generate_kwargs)
            
            return result
            
        except Exception as e:
            logger.error(f"转录失败: {str(e)}")
            raise
    
    def transcribe_batch(
        self,
        audio_list: List[Union[str, np.ndarray, Dict]],
        language: Optional[str] = None,
        task: str = "transcribe",
        batch_size: int = 8,
        **kwargs
    ) -> List[Dict]:
        """
        批量转录音频
        
        Args:
            audio_list: 音频文件路径或音频数组列表
            language: 语言代码（可选）
            task: 任务类型
            batch_size: 批处理大小
            **kwargs: 其他生成参数
        
        Returns:
            转录结果列表
        """
        try:
            # max_new_tokens 需要小于 max_target_positions (448)
            # 考虑到 decoder_input_ids 已经包含特殊令牌（通常为 4），设置为 444 以确保不超过限制
            generate_kwargs = {
                "max_new_tokens": 444,
                "num_beams": 1,
                "condition_on_prev_tokens": False,
                "compression_ratio_threshold": 1.35,
                "temperature": (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
                "logprob_threshold": -1.0,
                "no_speech_threshold": 0.6,
                **kwargs
            }
            
            if language:
                generate_kwargs["language"] = language
            
            if task == "translate":
                generate_kwargs["task"] = "translate"
            
            results = self.pipe(audio_list, batch_size=batch_size, generate_kwargs=generate_kwargs)
            
            return results if isinstance(results, list) else [results]
            
        except Exception as e:
            logger.error(f"批量转录失败: {str(e)}")
            raise


